import java.util.*;

public class Main {
    public static void main(String[] args) {
        Solution1 s = new Solution1();
        System.out.println(Arrays.toString(s.topKFrequent(new int[]{1,1,1,2,2,3}, 2)));
        System.out.println(Arrays.toString(s.topKFrequent(new int[]{1}, 1)));
    }
}


//HashMap统计词频，然后最大堆找最高频k个
class Solution {
    static class Num {
        int value;
        int count;
        Num(int value, int count) { this.value = value; this.count = count; }
    }

    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> counts = new HashMap<>();
        for (int num : nums) {
            counts.put(num, counts.getOrDefault(num, 0) + 1);
        }

        Queue<Num> minHeap = new PriorityQueue<>(Comparator.comparingInt(v -> v.count));
        for (Map.Entry<Integer, Integer> entry : counts.entrySet()) {
            if (minHeap.size() == k) {
                if (minHeap.peek().count < entry.getValue()) {
                    minHeap.poll();
                    minHeap.offer(new Num(entry.getKey(), entry.getValue()));
                }
            } else {
                minHeap.offer(new Num(entry.getKey(), entry.getValue()));
            }
        }

        int[] ans = new int[minHeap.size()];
        for (int i = 0; i < ans.length; i++) {
            ans[i] = minHeap.poll().value;
        }
        return ans;
    }
}

//类快速排序法
class Solution1 {
    static class Num {
        int value;
        int count;
        Num(int value, int count) { this.value = value; this.count = count; }
    }

    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> counts = new HashMap<>();
        for (int num : nums) {
            counts.put(num, counts.getOrDefault(num, 0) + 1);
        }

        Num[] allOccurrences = new Num[counts.size()];
        int i = 0;
        for (Map.Entry<Integer, Integer> entry : counts.entrySet()) {
            allOccurrences[i++] = new Num(entry.getKey(), entry.getValue());
        }

//        quickSelectRecursive(allOccurrences, 0, allOccurrences.length - 1, k);
        quickSelectLoop(allOccurrences, k);

        int start = allOccurrences.length - k;
        int[] ans = new int[k];
        for (i = 0; i < k; i++) {
            ans[i] = allOccurrences[start + i].value;
        }
        return ans;
    }

    //quickSelect: 返回最大的k个数，循环实现
    private int[] quickSelectLoop(Num[] s, int k) {
        int l = 0, h = s.length - 1;
        int target = s.length - k;
        while (true) {
            int mid = partition(s, l, h);
            if (mid == target) {
                int[] result = new int[k];
                for (int i = 0; i < k; i++) {
                    result[i] = s[mid + i].value;
                }
                return result;
            }
            if (mid < target) {
                l = mid + 1;
            } else {
                h = mid - 1;
            }
        }
    }

    //quickSelect: 返回最大的k个数，递归实现
    private int[] quickSelectRecursive(Num[] s, int l, int h, int k) {
        int mid = partition(s, l, h);
        int target = s.length - k;
        if (mid == target) {
            int[] result = new int[k];
            for (int i = 0; i < k; i++) {
                result[i] = s[mid + i].value;
            }
            return result;
        }
        if (mid < target) {
            return quickSelectRecursive(s, mid + 1, h, k);
        } else {
            return quickSelectRecursive(s, l, mid - 1, k);
        }
    }

    private int partition(Num[] s, int l, int h) {
        Num pivot = s[h];
        int i = l, j = l;
        while (j < h) {
            if (s[j].count < pivot.count) {
                swap(s, i++, j);
            }
            j++;
        }
        swap(s, i, h);
        return i;
    }

    private void swap(Num[] s, int i, int j) {
        Num temp = s[i];
        s[i] = s[j];
        s[j] = temp;
    }
}
